#include "kernel_operator.h"
#include <cstdint>
#include <cmath>

constexpr int32_t BUFFER_NUM = 2; // tensor num for each queue

// 定义复数类型
struct complex64 {
    float real;
    float imag;
};

// 定义复数类型
struct complex128 {
    float real_1;
    float real_2;
    float imag_1;
    float imag_2;
};

template<typename TYPE_X, typename TYPE_OUT> class KernelIsRealObj {
public:
    __aicore__ inline KernelIsRealObj() {}

    __aicore__ inline void Init(GM_ADDR x, GM_ADDR out, uint32_t smallCoreDataNum,
                                uint32_t bigCoreDataNum, uint32_t finalBigTileNum, 
                                uint32_t finalSmallTileNum, uint32_t tileDataNum, 
                                uint32_t smallTailDataNum, uint32_t bigTailDataNum, 
                                uint32_t tailBlockNum, uint32_t dataTypeId)
    {
        ASSERT(AscendC::GetBlockNum() != 0 && "block dim can not be zero!");
        uint32_t coreNum = AscendC::GetBlockIdx();
        uint32_t globalBufferIndex = bigCoreDataNum * AscendC::GetBlockIdx();
        this->dataTypeId = dataTypeId;

        this->tileDataNum = tileDataNum;
        if (coreNum < tailBlockNum) { 
          this->coreDataNum = bigCoreDataNum;
          this->tileNum = finalBigTileNum;
          this->tailDataNum = bigTailDataNum;
        }
        else { 
          this->coreDataNum = smallCoreDataNum;
          this->tileNum = finalSmallTileNum;
          this->tailDataNum = smallTailDataNum;
          globalBufferIndex -= (bigCoreDataNum - smallCoreDataNum) * (AscendC::GetBlockIdx() - tailBlockNum);
        }
        xGm.SetGlobalBuffer((__gm__ TYPE_X*)x + globalBufferIndex, this->coreDataNum);
        outGm.SetGlobalBuffer((__gm__ bool *)out, 1);
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileDataNum * sizeof(TYPE_X));
        pipe.InitBuffer(outQueueBool, BUFFER_NUM, 1);
    }

    __aicore__ inline void Process()
    {
        this->processDataNum = this->tileDataNum;
        CopyIn(0);
        Compute(0);
        CopyOut(0);
    }

private:
    __aicore__ inline void CopyIn(int32_t progress)
    {
      AscendC::LocalTensor<TYPE_X> xLocal = inQueueX.AllocTensor<TYPE_X>();
      AscendC::DataCopy(xLocal, xGm[progress * this->tileDataNum], this->processDataNum);
      inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void Compute(int32_t progress) {
        AscendC::LocalTensor<TYPE_X> xLocal = inQueueX.DeQue<TYPE_X>();
        AscendC::LocalTensor<TYPE_OUT> outLocal = outQueueBool.AllocTensor<TYPE_OUT>();
        
        bool isComplex = (dataTypeId == 16 || dataTypeId == 17);

        if (!isComplex) {
            outLocal.SetValue(0, 1);
        }
        else{
            outLocal.SetValue(0, 0);
        }

        outQueueBool.EnQue(outLocal);
        inQueueX.FreeTensor(xLocal);
    }

    __aicore__ inline void CopyOut(int32_t progress)
    {
      AscendC::LocalTensor<TYPE_OUT> outLocal = outQueueBool.DeQue<TYPE_OUT>();
      AscendC::DataCopyExtParams copyParams{1, (uint32_t)(1), 0, 0, 0};
      
      AscendC::DataCopyPad(outGm[0], outLocal, copyParams);
      
      outQueueBool.FreeTensor(outLocal);
    }
    
private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueBool;
    AscendC::GlobalTensor<TYPE_X> xGm;
    AscendC::GlobalTensor<TYPE_OUT> outGm;

    AscendC::DataCopyExtParams copyParams;

    uint32_t coreDataNum;
    uint32_t tileNum;
    uint32_t tileDataNum;
    uint32_t tailDataNum;
    uint32_t processDataNum;
    uint32_t dataTypeId;
};

extern "C" __global__ __aicore__ void is_real_obj_custom(GM_ADDR x, GM_ADDR out, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    KernelIsRealObj<DTYPE_X, DTYPE_OUT> op;
    op.Init(x, out, tiling_data.smallCoreDataNum, 
            tiling_data.bigCoreDataNum, tiling_data.finalBigTileNum, 
            tiling_data.finalSmallTileNum, tiling_data.tileDataNum, 
            tiling_data.smallTailDataNum, tiling_data.bigTailDataNum, 
            tiling_data.tailBlockNum,tiling_data.dataTypeId);  
    op.Process();
}
